#include <windows.h>
#include <stdio.h>
#include <tlhelp32.h>
#include "interface.h"
#include "overlay/overlay.h"

bool CheckProcess(DWORD pID)
{
  bool hasModule = false;
  HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, pID);
  if (hSnapshot != INVALID_HANDLE_VALUE)
  {
    MODULEENTRY32 me;
    me.dwSize = sizeof me;
    if (Module32First(hSnapshot, &me))
    {
      do
      {
        if (!wcsicmp(me.szModule, L"d3d9.dll") ||
            !wcsicmp(me.szModule, L"d3d10.dll") ||
            !wcsicmp(me.szModule, L"d3d11.dll"))
          hasModule = true;
      } while (Module32Next(hSnapshot, &me));
    }
    CloseHandle(hSnapshot);
  }
  return hasModule;
}
BOOL CALLBACK EnumWindowsCallback(HWND hWnd, LPARAM lParam)
{
  if (!IsWindowVisible(hWnd)) return TRUE;
  HWND hResponse = (HWND) lParam;
  DWORD pID;
  DWORD tID = GetWindowThreadProcessId(hWnd, &pID);
  if (CheckProcess(pID))
    PostMessage(hResponse, KAPPA64_RESPONSE_LISTPROC, (WPARAM) hWnd, pID);
  return TRUE;
}

int WINAPI wWinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPWSTR lpCmdLine, int nCmdShow)
{
  HWND hWndHost = (HWND) _wtoi(lpCmdLine);
  DWORD pHostID;
  GetWindowThreadProcessId(hWndHost, &pHostID);
  HANDLE hHandles[2];
  hHandles[0] = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, pHostID);
  if (hHandles[0] == NULL)
    return 0;
  MSG msg;
  PeekMessage(&msg, NULL, WM_USER, WM_USER, PM_NOREMOVE);

  HMODULE hDLL = GetModuleHandle(L"kappa64");
  hHandles[1] = NULL;
  HHOOK hHook = NULL;
  HHOOK hMouseHook = NULL;
  HHOOK hKeyboardHook = NULL;
  do
  {
    MsgWaitForMultipleObjects(hHandles[1] ? 2 : 1, hHandles, FALSE, INFINITE, QS_ALLEVENTS);

    if (PeekMessage(&msg, NULL, 0, 0, PM_REMOVE))
    {
      switch (msg.message)
      {
      case KAPPA64_REQUEST_LISTPROC:
        EnumWindows(EnumWindowsCallback, (LPARAM) hWndHost);
        break;
      case KAPPA64_REQUEST_HOOK:
        if (hMouseHook) UnhookWindowsHookEx(hMouseHook);
        if (hHook) UnhookWindowsHookEx(hHook);
        hMouseHook = NULL;
        hHook = SetWindowsHookEx(WH_CBT, HookProc, hDLL, msg.lParam);
        if (hHook)
          hHandles[1] = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, msg.wParam);
        else
          PostMessage(hWndHost, KAPPA64_RESPONSE_LOST, 0, 0);
        break;
      case KAPPA64_REQUEST_UNHOOK:
        if (hMouseHook) UnhookWindowsHookEx(hMouseHook);
        if (hHook) UnhookWindowsHookEx(hHook);
        hMouseHook = NULL;
        hHook = NULL;
        break;
      case KAPPA64_REQUEST_HOOKMOUSE:
        if (hMouseHook) UnhookWindowsHookEx(hMouseHook);
        if (hKeyboardHook) UnhookWindowsHookEx(hKeyboardHook);
        hMouseHook = SetWindowsHookEx(WH_MOUSE, MouseHookProc, hDLL, msg.lParam);
        hKeyboardHook = SetWindowsHookEx(WH_KEYBOARD, KeyboardHookProc, hDLL, msg.lParam);
        break;
      case KAPPA64_REQUEST_UNHOOKMOUSE:
        if (hMouseHook) UnhookWindowsHookEx(hMouseHook);
        if (hKeyboardHook) UnhookWindowsHookEx(hKeyboardHook);
        hMouseHook = NULL;
        hKeyboardHook = NULL;
        break;
      }
    }

    DWORD code;
    if (GetExitCodeProcess(hHandles[0], &code) && code != STILL_ACTIVE)
      break;
    if (hHandles[1] && GetExitCodeProcess(hHandles[1], &code) && code != STILL_ACTIVE)
    {
      if (hWndHost) PostMessage(hWndHost, KAPPA64_RESPONSE_LOST, 0, 0);
      hHandles[1] = NULL;
    }
  } while (msg.message != WM_QUIT);
  if (hMouseHook) UnhookWindowsHookEx(hMouseHook);
  if (hKeyboardHook) UnhookWindowsHookEx(hKeyboardHook);
  if (hHook) UnhookWindowsHookEx(hHook);
  if (hHandles[1]) CloseHandle(hHandles[1]);
  CloseHandle(hHandles[0]);
  return 0;
}
